#!/usr/bin/env python


import sys
import getopt
import time
import os
import subprocess

# "solver" should point to the binary built. 
# e.g. "/home/z3-str/str" or "/home/work/tool/z3/myStrTheory/str"
solver = "/Users/hellok/bctf/week/20140507"
# Before sovling the constraints, Z3-str has to properly encodes inputs as intermediated 
# format. (E.g. "\n" and "\t" should be encoded). Variable "tmpEncodingDir" where the intermediated
# inputs are saved
tmpEncodingDir = "/tmp/z3_str_convert"
clearTempFile = 1
#=================================================================== 

def encodeConstStr(constStr):
  constStr = constStr.replace(' ',    '_aScIi_040')
  constStr = constStr.replace('\\\"', '_aScIi_042')
  constStr = constStr.replace('#',    '_aScIi_043')
  constStr = constStr.replace('$',    '_aScIi_044')
  constStr = constStr.replace('.',    '_aScIi_046')
  constStr = constStr.replace('\'',   '_aScIi_047')
  constStr = constStr.replace('(',    '_aScIi_050')
  constStr = constStr.replace(')',    '_aScIi_051')
  constStr = constStr.replace(',',    '_aScIi_054')
  constStr = constStr.replace(':',    '_aScIi_072')
  constStr = constStr.replace(';',    '_aScIi_073')
  constStr = constStr.replace('[',    '_aScIi_133')
  constStr = constStr.replace(']',    '_aScIi_135')
  constStr = constStr.replace('\\\\', '_aScIi_134')
  constStr = constStr.replace('{',    '_aScIi_173')
  constStr = constStr.replace('}',    '_aScIi_175')
  constStr = constStr.replace('|',    '_aScIi_174')  
  constStr = constStr.replace('`',    '_aScIi_140')  
  constStr = constStr.replace('\\t',  '_aScIi_011')
  constStr = constStr.replace('\\n',  '_aScIi_012')
  return constStr


def encodeSolution(constStr):
  if constStr.find("\\\"") != -1:
    constStr = constStr.replace("\\\"", "\\\\\\\"")
  elif constStr.find("\\\'") != -1:
    constStr = constStr.replace("\\\"", "\\\\\\\'")
  elif constStr.find("\\\\") != -1:
    constStr = constStr.replace("\\\\", "\\\\\\\\")    
  elif constStr.find("\n") != -1:
    constStr = constStr.replace("\n", "\\n") 
  elif constStr.find("\t") != -1:
    constStr = constStr.replace("\t", "\\t")     
  elif constStr.find("\\t") != -1:
    constStr = constStr.replace("\\t", "\\\\t")      
  elif constStr.find("\\n") != -1:
    constStr = constStr.replace("\\n", "\\\\n")    
  elif constStr.find("\\") != -1:
    constStr = constStr.replace("\\", "\\\\")     
  return constStr



varTypeDict = {}
varSolution = {}
fileContent = ""
iFile = ""



def convert(org_file, convertedOriginalFile):  
  global varTypeDict
  global fileContent
  global iFile  
  
  f_o = open(org_file, 'r')  
  f_n = open(convertedOriginalFile, 'w')

  declared_string_var = []
  declared_string_const = []
  converted_cstr = ""
  
  fileContent = f_o.read()  
  f_o.close()
  linesInFile = fileContent.split("\n")
  
  for line in linesInFile:    
    line = line.strip();
    line = line.replace('\t', ' ')
    if line == "":
      continue
    elif line.startswith(';'):
      continue
    elif line.startswith('%'):
      continue
    elif line.startswith('//'):
      continue
    elif line.find("get-model") != -1:
      continue
    elif line.find("set-option") != -1:
      continue
    elif line.find("declare-variable") != -1:    
      declared_string_var.append(line.replace('declare-variable', 'declare-const'))
      # collect types of variables
      tt = line[1:-1].lstrip()
      tt = tt[tt.find("declare-variable") + 16:].strip()      
      varName = tt[0:tt.find(' ')].strip()
      varType = tt[tt.find(' ') + 1: ].strip().lower()
      varTypeDict[varName] = varType      
      continue 
    
    elif line.find("declare-const") != -1:
      # collect types of variables
      tt = line[1:-1].lstrip()
      tt = tt[tt.find("declare-const") + 13:].strip()
      varName = tt[0:tt.find(' ')].strip()
      varType = tt[tt.find(' ') + 1: ].strip().lower()
      varTypeDict[varName] = varType
    
    elif line.find("declare-fun") != -1 :
      # collect types of variables
      tt = line[1:-1].lstrip()
      tt = tt[tt.find("declare-fun") + 11:].strip()
      tt = tt[:tt.find('(')] + tt[tt.find(')') + 1:]
      varName = tt[0:tt.find(' ')].strip()
      varType = tt[tt.find(' ') + 1: ].strip().lower()
      varTypeDict[varName] = varType      
    
    # ------------------------------------
    # start: processing const string
    p1 = -1
    while True:
      p1 = line.find('\"', p1 + 1);
      if p1 == -1:
        break;
      
      # exclude the case "str\"str\"str"
      p2 = p1 + 1
      while p2 < len(line):
        if line[p2] == "\\":
          p2 = p2 + 2
          continue
        elif line[p2] == "\"":
          break
        p2 = p2 + 1

        if p2 >= len(line):
          p2 = -1
    
      if p2 == -1:
        print('input format error!\n')
        print line
        return "eRrOr"
      
      
      old_s = line[p1: p2 + 1]
      encoded_s = encodeConstStr( old_s[1 : len(old_s) - 1] )
      line = line.replace(old_s, '__cOnStStR_' + encoded_s)
      
      if encoded_s not in declared_string_const:
        declared_string_const.append(encoded_s)        
      p1 = p2
    # -----------------------------
    # end: processing const string
    converted_cstr = converted_cstr + line + '\n'
    
  output_str = ""
  for strv in declared_string_var:
    output_str = output_str + strv + "\n"  
  output_str = output_str + '\n'  
  for str_const in declared_string_const:
    output_str = output_str + '(declare-const __cOnStStR_' + str_const + ' String)\n'    
  output_str = output_str + '\n'
  output_str = output_str + converted_cstr    
  f_n.write(output_str)        
  f_n.close()  
  

  

def processOutput(output):
  global varSolution
  
  if output.find("(error \"line") >= 0:
    res = ''
    lines = output.split("\n")
    for line in lines:
      if line.startswith('(error '):
	res = res + line + "\n"      
    return res
  
  output = output.replace("\t", "\\t")  
  lines = output.split("\n")
  result = ""
  oneItem = ""
  for line in lines:
    line = line.lstrip(' ');    
    line = line.replace("\n", "\\n")
    if line.strip() == "":
      continue
    elif line.startswith('_t_'):
      continue
    elif line.startswith('unique-value!'):
      continue
    
    elif line.startswith('**************') or line.startswith('>> ') or line.startswith('--------------'):
      if oneItem != "":
        result = result + oneItem + "\n"
        fields = oneItem.split(" -> ")
        varName = fields[0]
        varValue = fields[len(fields) - 1]      

        varSolution[varName] = encodeSolution(varValue)
        
      result = result + line + "\n"     
      continue    
    
    elif line.find(" -> ") == -1:
      oneItem = oneItem + "\n" + line
      #result = result[:-1] + "\\n" + line + "\n"
      continue
    
    elif line.find(" -> ") != -1:
      if oneItem != "":
        result = result + oneItem + "\n"        
        fields = oneItem.split(" -> ")
        varName = fields[0]
        varValue = fields[len(fields) - 1]
        varSolution[varName] = encodeSolution(varValue)
      oneItem = line
  return result




def genSolAsserts():
  global varTypeDict
  global varSolution
  result = ""
  for name, value in varSolution.items():
    if (varTypeDict[name] == "string" and value != '""'):
      result = result + "(assert (= " + name + " \"" + value + "\"))" + "\n"
    elif varTypeDict[name] == "bool":
      if value == "true":
        result = result + "(assert " + name + ")" + "\n"
      else:
        result = result + "(assert (not " + name + "))" + "\n"
    else:
      result = result + "(assert (= " + name + " " + value + "))" + "\n"
  return result
  
  


def printUseage():
  print 'USAGE: '
  print '  Z3-str.py -f <inputfile>\n'  
  print '\n'
  
      

if __name__ == '__main__':
  argv = sys.argv[1:]
  inputFile = '';  
  allowLoopCut = 0
  modifyFreeLen = 0
  freeVarMaxLen = 0
  
  try:
    opts, args = getopt.getopt(argv,"hf:")
  except getopt.GetoptError:
    printUseage()
    sys.exit()
  for opt, arg in opts:
    if opt == '-h':      
      printUseage()
      sys.exit()
    elif opt in ("-f"):
      inputFile = arg
    
      
  if inputFile == '':
    printUseage()
    sys.exit()
  
  if not os.path.exists(inputFile):
    print "Error: Input file does not exist: \"" + inputFile + "\""    
    exit(0)
   
  convertDir = tmpEncodingDir;
  if not os.path.exists(convertDir):
    os.makedirs(convertDir)
  
  fileName = os.path.basename(inputFile);    
  iFile = fileName
  convertedOriginalFile = os.path.join(convertDir, fileName)   
  
  rr = convert(inputFile, convertedOriginalFile)
  if rr == "eRrOr":
    exit(0)
    
  try:
    start = time.time() 
    paras = [solver, "-f", convertedOriginalFile]
     
    # --------------------------------------------------  
    # Solve the origial constraints.
    #   > SAT, verify the ansert
    #   > UNSAT/UNKNOWN, report
    # --------------------------------------------------
    err = subprocess.check_output(paras, );
    outStr = processOutput(err)
    
    # Remove the temp file of the encoding of the original input
    if clearTempFile == 1:
      if os.path.exists(convertedOriginalFile):
        os.remove(convertedOriginalFile)
    
    # --------------------------------------------------
    # Result parsing
    #  - If SAT, assert the solutions as additional constraints to verify the answer
    #    ** Verificaiton OK:   print "* v-ok"
    #    ** Verification FAIL: print "* v-fail" together with "UNKNOWN"
    #  - If UNSAT/UNKNOWN
    # --------------------------------------------------
    if outStr.find("String!val!") != -1:
        print "* v-fail"
        print "************************"
        print ">> UNKNOWN  (1)"
        print "************************" 
    elif outStr.find(">> SAT") != -1: 
      ii = fileContent.find("(check-sat)")    
      solutionAssert = genSolAsserts()
      newInputStr = fileContent[ :ii - 1] + "\n\n\n" + solutionAssert + "\n\n" + fileContent[ii:]
      verifyInputFilename = tmpEncodingDir + "/verify_" + iFile
      convertedVerifyInputFilename = tmpEncodingDir + "/verify_conv_" + iFile 
      f_n = open(verifyInputFilename, 'w')
      f_n.write(newInputStr)        
      f_n.close()
      
      convert(verifyInputFilename, convertedVerifyInputFilename)
      paras = [solver, "-f", convertedVerifyInputFilename]
      
      # run the solver again, check the solution.
      err1 = subprocess.check_output(paras, );
      if err1.find(">> SAT") != -1:
        print "* v-ok" 
        sys.stdout.write(outStr)  
        eclapse = (time.time() - start)
      else:
        print "* v-fail"
        print "************************"
        print ">> UNKNOWN (2)"
        print "************************"
      if clearTempFile == 1:
        if os.path.exists(verifyInputFilename):
          os.remove(verifyInputFilename)
        if os.path.exists(convertedVerifyInputFilename):
          os.remove(convertedVerifyInputFilename)
    else:
      sys.stdout.write(outStr)  
      eclapse = (time.time() - start)
  except KeyboardInterrupt:
    print "Interrupted by keyborad"
    
  
